import asyncio
import functools
import time
from typing import Any, AsyncGenerator, Callable, Generator, TypeVar, Union

from logger_config_db import setup_logger

logger = setup_logger(__name__)

F = TypeVar("F", bound=Callable[..., Any])


def timeit(func: F) -> F:
    """
    计算函数执行时间的装饰器
    支持:
        1. 普通函数
        2. 异步函数
        3. 生成器函数
        4. 异步生成器函数
    """

    @functools.wraps(func)
    def async_wrapper(*args, **kwargs) -> Any:
        start_time = time.perf_counter()
        try:
            result = func(*args, **kwargs)

            # 处理生成器函数
            if isinstance(result, (Generator, AsyncGenerator)):
                return timeit_generator(result, start_time, func.__name__)

            end_time = time.perf_counter()
            execution_time = end_time - start_time
            logger.info(
                f"Function '{func.__name__}' took {execution_time:.4f} seconds to execute"
            )
            return result

        except Exception as e:
            end_time = time.perf_counter()
            execution_time = end_time - start_time
            logger.error(
                f"Function '{func.__name__}' failed after {execution_time:.4f} seconds with error: {str(e)}"
            )
            raise

    async def async_wrapper(*args, **kwargs) -> Any:
        start_time = time.perf_counter()
        try:
            result = await func(*args, **kwargs)

            # 处理异步生成器
            if isinstance(result, AsyncGenerator):
                return timeit_async_generator(result, start_time,
                                              func.__name__)

            end_time = time.perf_counter()
            execution_time = end_time - start_time
            logger.info(
                f"Async function '{func.__name__}' took {execution_time:.4f} seconds to execute"
            )
            return result

        except Exception as e:
            end_time = time.perf_counter()
            execution_time = end_time - start_time
            logger.error(
                f"Async function '{func.__name__}' failed after {execution_time:.4f} seconds with error: {str(e)}"
            )
            raise

    async def timeit_async_generator(gen, start_time, func_name):
        try:
            async for item in gen:
                yield item
        finally:
            end_time = time.perf_counter()
            execution_time = end_time - start_time
            logger.info(
                f"Async generator '{func_name}' took {execution_time:.4f} seconds to complete"
            )

    def timeit_generator(gen, start_time, func_name):
        try:
            yield from gen
        finally:
            end_time = time.perf_counter()
            execution_time = end_time - start_time
            logger.info(
                f"Generator '{func_name}' took {execution_time:.4f} seconds to complete"
            )

    # 根据函数类型选择合适的装饰器
    if asyncio.iscoroutinefunction(func):
        return async_wrapper
    return async_wrapper
